import os

#os.environ['QIANFAN_ACCESS_KEY']='**********'
#os.environ['QIANFAN_SECRET_KEY']='**********'
#os.environ['BAIDU_MODEL']='ERNIE-Speed-128K'

if 'BAIDU_MODEL' not in os.environ:
    os.environ['BAIDU_MODEL'] = 'ERNIE-Speed-128K'

import qianfan
import time
import json
from typing import List, Optional

import logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

client = qianfan.ChatCompletion()

class ResponseFormatError(Exception):
    pass

def _get_json_str(response: str) -> str:
    start_index = response.find("```json")
    if start_index < 0:
        raise ResponseFormatError("答案生成失败，找不到起始“```json”。下次我注意严格按照要求生成Markdown合法的JSON对象。")

    end_index = response.find("```", start_index + 8)
    if end_index < 0:
        raise ResponseFormatError("答案生成失败，找不到结束“```”。下次我注意严格按照要求生成Markdown合法的JSON对象。")

    return response[start_index + 8:end_index]

def _parse_messages(messages: List[str]) -> dict:
    system = None
    for item in messages:
        if 'system' == item['role']:
            system = item['content']
            break

    messages_return = []
    for item in messages:
        if 'system' == item['role']:
            continue
        messages_return.append({"role": item['role'], "content": item['content']})

    return {
        "messages": messages_return,
        "system": system
    }

def _get_response(messages: List[dict], max_tokens: Optional[int]=None) -> dict:
    messages_for_qianfan = _parse_messages(messages)
    response = client.do(
        model=os.environ['BAIDU_MODEL'],
        messages=messages_for_qianfan['messages'],
        system=messages_for_qianfan['system'],
        stream=False,
        max_output_tokens=max_tokens
    )
    return response

def make_api_call(messages, max_tokens: Optional[int]=None, is_final_answer=False, custom_client=None, max_attempts: int = 3):
    logger.info('Making API call with messages: %s (max_tokens=%s,is_final_answer=%s)'%(json.dumps(messages), str(max_tokens), str(is_final_answer)))
    global client
    if custom_client != None:
        client = custom_client

    for attempt in range(max_attempts):
        try:
            if is_final_answer:
                response = _get_response(messages, max_tokens)
                logger.debug(f'Response: {response["result"]}')
                return response['result']

            else:
                try:
                    send_messages = messages
                    for json_err_time in range(3):
                        try:
                            response = _get_response(send_messages, max_tokens)
                            json_str = _get_json_str(response['result'])
                            break
                        except ResponseFormatError:
                            logger.debug('Error parsing JSON from response. Retrying...')
                            if json_err_time == 0:
                                send_messages[-1]['content'] += '''


示例：

```json
{
    "title": "识别关键信息",
    "content": "为了解决这个问题，我们需要仔细检查给定的信息，识别关键元素，这些元素将指导我们的解决方案过程。这涉及到了解...",
    "next_action": "continue"
}```

你**必须**按照这个格式返回一个JSON对象。'''
                            time.sleep(1)
                except Exception as json_error:
                    logger.warning(f'Error parsing JSON from response: {str(json_error)}. Response: {response["result"]}')
                    raise json_error

                json_str = json_str.replace('\n', ' ')
                json_str = json_str.replace('\r', ' ')
                logger.debug(f'Response JSON: {json_str}')
                result = json.loads(json_str)
                result = result[0] if isinstance(result, list) else result

                if not isinstance(result, dict):
                    raise Exception('Invalid JSON from response')

                return result

        except Exception as e:
            logger.debug('Exception type: %s. Exception message: %s'%(type(e), str(e)))
            if attempt == max_attempts - 1:
                if is_final_answer:
                    exception_return = {"title": "Error", "content": f"Failed to generate final answer after 3 attempts. Error: {str(e)}"}
                else:
                    exception_return = {"title": "Error", "content": f"Failed to generate step after 3 attempts. Error: {str(e)}", "next_action": "final_answer"}
                logger.error(f'Exception: {str(e)}. Returning exception_return: {exception_return}')
                return exception_return
            time.sleep(1)  # Wait for 1 second before retrying
        logger.info('Retrying API call...')

def generate_response(prompt: str, custom_client=None):
    logger.info(f'Generating response for prompt: {prompt}')

    system_content = '''你是一个专家级的人工智能助手，能够逐步解释你的推理过程。对于每一步，提供一个标题来描述你在那一步正在做什么，以及内容。决定是否需要进一步的步骤，或者是否准备好给出最终答案。以JSON格式响应，包含'title'、'content'和'next_action'（可以是'continue'或'final_answer'）键。尽可能使用多的推理步骤。至少3步。意识到你作为语言模型的局限性，你能做什么和不能做什么。在你的推理中，包括对替代答案的探索。考虑你可能是错的，并且如果你的推理错了，错在哪里。完全测试所有其他可能性。你可以错。当你你说你正在重新检查时，实际上要重新检查，并使用另一种方法来做。不要只是说你在重新检查。使用至少3种方法得出答案。使用最佳实践。

有效JSON响应的示例：

```json
{
    "title": "识别关键信息",
    "content": "为了解决这个问题，我们需要仔细检查给定的信息，识别关键元素，这些元素将指导我们的解决方案过程。这涉及到了解...",
    "next_action": "continue"
}```

你**必须**返回一个JSON对象而不是列表。如果你要返回多条内容你要把内容通过常规的文本格式放在`content`属性里。'''

    messages = [
        {"role": "system", "content": system_content},
        {"role": "user", "content": prompt}
    ]

    steps = []
    step_count = 1
    total_thinking_time = 0

    while True:
        start_time = time.time()
        step_data = make_api_call(messages, None, custom_client=custom_client)
        end_time = time.time()
        thinking_time = end_time - start_time
        total_thinking_time += thinking_time

        steps.append((f"Step {step_count}: {step_data['title']}", step_data['content'], thinking_time))

        json_str = json.dumps(step_data, ensure_ascii=False, indent=4)
        messages.append({'role': 'assistant', 'content': f'```json\n{json_str}\n```'})

        if step_data['next_action'] == 'final_answer' or step_count > 25: # Maximum of 25 steps to prevent infinite thinking time. Can be adjusted.
            break
        messages.append({"role": "user", "content": "请继续。"})

        step_count += 1

        # Yield after each step for Streamlit to update
        yield steps, None  # We're not yielding the total time until the end

    # Generate final answer
    messages.append({"role": "user", "content": "请仅基于上述推理提供最终答案。不要使用JSON格式。只提供文本响应，不要有任何标题或前言。保留原始提示所指示的任何格式，例如自由回答或多选题的精确格式。"})

    start_time = time.time()
    final_data = make_api_call(messages, None, is_final_answer=True, custom_client=custom_client)
    end_time = time.time()
    thinking_time = end_time - start_time
    total_thinking_time += thinking_time

    steps.append(("Final Answer", final_data, thinking_time))

    yield steps, total_thinking_time
